#include <onnxruntime_cxx_api.h>
#include <iostream>
#include <vector>
#include <string>
#include <array>
#include <chrono>
#include <memory>

class OnnxEmbedder
{
public:
    OnnxEmbedder(const std::string &model_path, bool use_cuda = false);

    // return tensor
    std::vector<Ort::Value> Run(const std::vector<int64_t> &input_ids,
                                const std::vector<int64_t> &attention_mask,
                                const std::vector<int64_t> &token_type_ids);

    // return embedding vector
    std::vector<float> RunEmbdding(const std::vector<int64_t> &input_ids,
                                   const std::vector<int64_t> &attention_mask,
                                   const std::vector<int64_t> &token_type_ids);
    double Benchmark(const std::vector<int64_t> &input_ids,
                     const std::vector<int64_t> &attention_mask,
                     const std::vector<int64_t> &token_type_ids,
                     int num_runs = 100);

private:
    Ort::Env env_;
    Ort::SessionOptions session_options_;
    std::unique_ptr<Ort::Session> session_;
    std::unique_ptr<Ort::AllocatorWithDefaultOptions> allocator_;
    std::vector<std::string> input_names_str_;
    std::vector<std::string> output_name_str_;
};